use burn::{
    config::Config,
    data::dataloader::DataLoaderBuilder,
    module::Module,
    optim::AdamConfig,
    record::CompactRecorder,
    tensor::backend::AutodiffBackend,
    train::{
        LearnerBuilder,
        metric::{
            AccuracyMetric, CpuMemory, CpuTemperature, CpuUse, CudaMetric, IterationSpeedMetric,
            LossMetric,
        },
    },
};
use serde::{Deserialize, Serialize};

use crate::{
    data::{
        batcher::HWDBBatcher,
        dataset::{HWDBDatasetConfig, HWDBHdf5DatasetConfig},
    },
    model::basic::BasicModelConfig,
};

#[derive(Debug, Clone, Deserialize, Serialize)]
pub enum HWDBDatasetConfType {
    Folder(HWDBDatasetConfig),
    Hdf5(HWDBHdf5DatasetConfig),
}

#[derive(Config)]
pub struct TrainingConfig {
    pub dataset: HWDBDatasetConfType,
    pub model: BasicModelConfig,
    pub optimizer: AdamConfig,
    #[config(default = 10)]
    pub num_epoches: usize,
    #[config(default = 1024)]
    pub batch_size: usize,
    #[config(default = 4)]
    pub num_workers: usize,
    #[config(default = 1007)]
    pub seed: u64,
    #[config(default = 1e-4)]
    pub lr: f64,
    #[config(default = 184)]
    pub image_size: usize,
}

fn create_artifact_dir(artifact_dir: &str) {
    // Remove existing artifacts before to get an accurate learner summary
    std::fs::remove_dir_all(artifact_dir).ok();
    std::fs::create_dir_all(artifact_dir).ok();
}

pub fn train<B: AutodiffBackend>(artifact_dir: &str, config: TrainingConfig, device: B::Device) {
    create_artifact_dir(artifact_dir);
    config
        .save(format!("{artifact_dir}/config.json"))
        .expect("Config should be saved successfully");

    B::seed(config.seed);

    let batcher = HWDBBatcher::new(false, 184, 184, Some((0.9398248, 0.47693425, 269, 259)));

    let (dataloader_train, dataloader_test) = match config.dataset {
        HWDBDatasetConfType::Folder(hwdb_dataset_config) => {
            let hwdbdataset = hwdb_dataset_config.init();
            let dataloader_train = DataLoaderBuilder::new(batcher.clone())
                .batch_size(config.batch_size)
                .shuffle(config.seed)
                .num_workers(config.num_workers)
                .build(hwdbdataset.train);

            let dataloader_test = DataLoaderBuilder::new(batcher)
                .batch_size(config.batch_size)
                .shuffle(config.seed)
                .num_workers(config.num_workers)
                .build(hwdbdataset.test);
            (dataloader_train, dataloader_test)
        }
        HWDBDatasetConfType::Hdf5(hwdb_hdf5_ds_config) => {
            let dataloader_train = DataLoaderBuilder::new(batcher.clone())
                .batch_size(config.batch_size)
                .shuffle(config.seed)
                .num_workers(config.num_workers)
                .build(hwdb_hdf5_ds_config.train());

            let dataloader_test = DataLoaderBuilder::new(batcher)
                .batch_size(config.batch_size)
                .shuffle(config.seed)
                .num_workers(config.num_workers)
                .build(hwdb_hdf5_ds_config.test());
            (dataloader_train, dataloader_test)
        }
    };

    let learner = LearnerBuilder::new(artifact_dir)
        .metric_train_numeric(AccuracyMetric::new())
        .metric_valid_numeric(AccuracyMetric::new())
        .metric_train_numeric(LossMetric::new())
        .metric_valid_numeric(LossMetric::new())
        .metric_train_numeric(IterationSpeedMetric::new())
        .metric_train_numeric(CpuMemory::default())
        .metric_train_numeric(CpuTemperature::default())
        .metric_train_numeric(CpuUse::default())
        .metric_train(CudaMetric::default())
        .with_file_checkpointer(CompactRecorder::new())
        .devices(vec![device.clone()])
        .num_epochs(config.num_epoches)
        .summary()
        .build(
            config.model.init::<B>(&device),
            config.optimizer.init(),
            config.lr,
        );

    let model_trained = learner.fit(dataloader_train, dataloader_test);

    model_trained
        .save_file(format!("{artifact_dir}/model"), &CompactRecorder::new())
        .expect("Trained model should be saved successfully");
}
